import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, classification_report
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.ensemble import AdaBoostClassifier
import warnings
warnings.filterwarnings('ignore')


def data_scale(df):
    # 选取特征
    features = df.loc[:, ['峰度', '频谱熵', '带宽', '频率峰值间距', '偏度', '频谱峰值数']]
    # 创建标准化器
    scaler = StandardScaler()
    # 标准化数据
    features_scaled = pd.DataFrame(scaler.fit_transform(features), index=features.index, columns=features.columns)
    return features_scaled


def data_split(features, labels):
    # 随机划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(features, labels, test_size=0.3, random_state=42)

    return X_train, X_test, y_train, y_test


def evaluate_performance(y_true, y_pred, labels=None):
    # 混淆矩阵
    cm = confusion_matrix(y_true, y_pred, labels=labels)
    cm_df = pd.DataFrame(cm, index=labels, columns=labels)
    print("混淆矩阵:")
    print(cm_df)

    # 分类评估报告
    print("\n分类评估报告:")
    report = classification_report(y_true, y_pred, target_names=[str(label) for label in labels])
    print(report)


def KNN(X_train, X_test, y_train, y_test):
    knn = KNeighborsClassifier(n_neighbors=3)
    knn.fit(X_train, y_train)
    y_pred_knn = knn.predict(X_test)
    print('KNN')
    evaluate_performance(y_test, y_pred_knn, labels)


def SVM(X_train, X_test, y_train, y_test):
    svm = SVC(kernel='linear')
    svm.fit(X_train, y_train)
    y_pred_svm = svm.predict(X_test)
    print('SVM')
    evaluate_performance(y_test, y_pred_svm, labels)


def AdaBoost(X_train, X_test, y_train, y_test):
    ada_clf = AdaBoostClassifier(n_estimators=50, random_state=42)
    ada_clf.fit(X_train, y_train)
    y_pred_ada = ada_clf.predict(X_test)
    print('AdaBoost')
    evaluate_performance(y_test, y_pred_ada, labels)


if __name__ == '__main__':
    material_list = ['材料1', '材料2', '材料3', '材料4']
    labels = [1, 2, 3]

    for i in range(len(material_list)):
        df = pd.read_excel('../data/Q1_data.xlsx', sheet_name=material_list[i])
        tag = df['励磁波形']

        # 标准化数据集
        scaled_df = data_scale(df)

        # 划分数据集
        X_train, X_test, y_train, y_test = data_split(scaled_df, tag)

        print(f'对{material_list[i]}分类')
        # KNN
        KNN(X_train, X_test, y_train, y_test)

        # SVM
        SVM(X_train, X_test, y_train, y_test)

        # AdaBoost
        AdaBoost(X_train, X_test, y_train, y_test)
